# 加载预训练模型
from transformers import BertModel, BertTokenizer
import torch

# 定义训练设备
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

pretrained = BertModel.from_pretrained(r"D:\AI\HuggingFace\my-model-cache\bert-base-chinese").to(DEVICE)
print(pretrained.embeddings.word_embeddings)

# 定义下游任务模型，将主干网络所提取到的特征进行分类，这里是一个简单的二分类任务，因此使用了一个全连接层
class Model(torch.nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.fc = torch.nn.Linear(768, 2)
    def forward(self, input_ids, attention_mask, token_type_ids):
        # 上游任务不参与训练
        with torch.no_grad():
            out = pretrained(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids) 

        # 下游任务参与训练
        out = self.fc(out.last_hidden_state[:, 0]) # 取出[CLS]对应的向量
        out = torch.nn.functional.softmax(out, dim=1) # 进行softmax分类
        return out